"""
3D超高层建筑疏散仿真项目（实时可视化版）
包含三维楼梯动态、各楼层平面图、实时统计显示
"""

import numpy as np
import heapq
import math
import random
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation

# ==================== 参数配置 ====================
CELL_SIZE = 1.0       # 元胞尺寸（米）
FLOORS = 10           # 总楼层数
HEIGHT, WIDTH = 30, 30 # 每层尺寸
AGENTS_PER_FLOOR = 40  # 每层初始人数
SPEED_MEAN = 1.0       # 初始速度均值（米/秒）
SPEED_STD = 0.1        # 速度标准差
STAIR_SPEED_FACTOR = 0.7 # 楼梯速度系数
FLOOR_HEIGHT = 3.0     # 层高（米）
MAX_STEPS = 1000       # 最大仿真步数

# ==================== 环境生成 ====================
def build_environment():
    # 初始化三维环境数组（1表示可通行，0表示障碍，2表示楼梯）
    env = np.ones((FLOORS, HEIGHT, WIDTH), dtype=int)
    
    # 随机生成4个楼梯位置（所有楼层共用）
    random.seed(42)
    global STAIR_POSITIONS
    STAIR_POSITIONS = random.sample(
        [(y, x) for y in range(2, HEIGHT-2) for x in range(2, WIDTH-2)],
        k=4
    )
    
    # 生成各楼层布局
    for z in range(FLOORS):
        grid = np.zeros((HEIGHT, WIDTH), dtype=int)
        pattern_type = z % 4  # 四种不同布局模式
        
        # 模式0：十字交叉走廊
        if pattern_type == 0:
            grid[HEIGHT//2-1:HEIGHT//2+2, :] = 1
            grid[:, WIDTH//2-1:WIDTH//2+2] = 1
            
        # 模式1：对角线通道
        elif pattern_type == 1:
            for i in range(HEIGHT):
                grid[i, i] = 1
                grid[i, WIDTH-1-i] = 1
                
        # 模式2：环形走廊
        elif pattern_type == 2:
            grid[2:-2, 2] = 1
            grid[2:-2, WIDTH-3] = 1
            grid[2, 2:-2] = 1
            grid[HEIGHT-3, 2:-2] = 1
            
        # 模式3：网格状走廊
        else:
            grid[5::10, :] = 1
            grid[:, 5::10] = 1
        
        # 添加随机障碍物
        np.random.seed(z)
        for _ in range(5):
            ry = np.random.randint(1, HEIGHT-4)
            rx = np.random.randint(1, WIDTH-4)
            grid[ry:ry+3, rx:rx+3] = 0
            
        # 设置楼梯位置
        for (sy, sx) in STAIR_POSITIONS:
            grid[sy, sx] = 2
            
        # 生成最终楼层布局（翻转0/1以符合直觉：0=障碍，1=可行走）
        env[z] = np.where(grid == 0, 0, np.where(grid == 2, 2, 1))
    
    return env

# ==================== 静态势场计算 ====================
def compute_static_field(env):
    field = np.full(env.shape, np.inf)
    heap = []
    visited = np.zeros(env.shape, bool)
    
    # 初始化出口位置（地面层的楼梯和角落）
    EXIT_COORDS = STAIR_POSITIONS + [(0, 0), (HEIGHT-1, WIDTH-1)]
    EXIT_POSITIONS = [(0, y, x) for (y, x) in EXIT_COORDS]
    
    # 多源Dijkstra算法
    for pos in EXIT_POSITIONS:
        field[pos] = 0
        heapq.heappush(heap, (0.0, pos))
    
    # 26邻域方向（允许三维移动）
    directions = [(dz, dy, dx) 
                 for dz in (-1,0,1) 
                 for dy in (-1,0,1) 
                 for dx in (-1,0,1) 
                 if not (dz==0 and dy==0 and dx==0)]
    
    while heap:
        current_dist, (z, y, x) = heapq.heappop(heap)
        if visited[z, y, x]:
            continue
        visited[z, y, x] = True
        
        for dz, dy, dx in directions:
            nz, ny, nx = z+dz, y+dy, x+dx
            if 0 <= nz < FLOORS and 0 <= ny < HEIGHT and 0 <= nx < WIDTH:
                if env[nz, ny, nx] != 0:  # 跳过障碍物
                    # 计算移动成本（关键修复点）
                    step = CELL_SIZE * (
                        math.sqrt(1 + (FLOOR_HEIGHT/CELL_SIZE)**2 
                        if dz != 0  # 跨楼层移动
                        else CELL_SIZE  # 同层移动
                    )
                    new_dist = current_dist + step
                    
                    if new_dist < field[nz, ny, nx]:
                        field[nz, ny, nx] = new_dist
                        heapq.heappush(heap, (new_dist, (nz, ny, nx))
    
    return field

# ==================== 人员类定义 ====================
class Evacuee:
    def __init__(self, start_z, start_y, start_x):
        self.pos = (start_z, start_y, start_x)
        self.base_speed = random.gauss(SPEED_MEAN, SPEED_STD)
        self.distance_traveled = 0.0
        self.time_elapsed = 0.0
        self.step_counter = 0.0
        
    def get_speed(self, vertical_move):
        """计算当前速度（考虑体力消耗）"""
        fatigue_factor = 1 - 0.01 * (self.distance_traveled**0.71) / (
            1 - 0.01 * (self.time_elapsed**0.51) * math.exp(0.01*self.time_elapsed))
        speed = self.base_speed * fatigue_factor
        return speed * (STAIR_SPEED_FACTOR if vertical_move else 1.0)

# ==================== 初始化人员 ====================  
def initialize_agents(environment):
    agents = []
    for z in range(FLOORS):
        # 获取本层可行走区域
        valid_positions = list(zip(*np.where(environment[z] == 1)))
        random.shuffle(valid_positions)
        
        # 在本层随机位置生成人员
        for y, x in valid_positions[:AGENTS_PER_FLOOR]:
            agents.append(Evacuee(z, y, x))
    return agents

# ==================== 仿真主循环 ====================
def run_simulation(env, static_field):
    agents = initialize_agents(env)
    history = []
    occupied = set(a.pos for a in agents)
    
    for _ in range(MAX_STEPS):
        current_positions = []
        new_occupied = set()
        
        for agent in list(agents):  # 创建副本以便安全删除
            # 更新个人状态
            agent.time_elapsed += 1
            agent.step_counter += 1
            
            z, y, x = agent.pos
            required_steps = CELL_SIZE / agent.get_speed(False)
            
            if agent.step_counter >= required_steps:
                # 寻找可行移动方向
                possible_moves = []
                for dz in (-1, 0, 1):
                    for dy in (-1, 0, 1):
                        for dx in (-1, 0, 1):
                            if dz == dy == dx == 0:
                                continue
                                
                            nz = z + dz
                            ny = y + dy
                            nx = x + dx
                            
                            if (0 <= nz < FLOORS and 
                                0 <= ny < HEIGHT and 
                                0 <= nx < WIDTH and 
                                env[nz, ny, nx] != 0 and 
                                (nz, ny, nx) not in occupied and 
                                static_field[nz, ny, nx] <= static_field[z, y, x]):
                                
                                possible_moves.append((
                                    static_field[nz, ny, nx],
                                    (dz, dy, dx)
                                ))
                
                if possible_moves:
                    # 选择最优移动方向
                    min_value = min(m[0] for m in possible_moves)
                    best_moves = [m[1] for m in possible_moves if m[0] == min_value]
                    dz, dy, dx = random.choice(best_moves)
                    
                    # 执行移动
                    new_pos = (z+dz, y+dy, x+dx)
                    occupied.remove(agent.pos)
                    new_occupied.add(new_pos)
                    agent.pos = new_pos
                    
                    # 更新移动距离
                    move_dist = CELL_SIZE * (
                        math.sqrt(1 + (FLOOR_HEIGHT/CELL_SIZE)**2) 
                        if dz != 0 
                        else CELL_SIZE
                    )
                    agent.distance_traveled += move_dist
                    agent.step_counter = 0
                    
                    # 检查是否到达出口
                    if new_pos in EXIT_POSITIONS:
                        agents.remove(agent)
                
                else:  # 无法移动
                    agent.step_counter = 0
                    new_occupied.add(agent.pos)
            
            current_positions.append(agent.pos)
        
        # 更新占位信息
        occupied = new_occupied
        
        # 记录当前状态
        remaining = len(agents)
        floor_distribution = [sum(1 for p in current_positions if p[0] == z) 
                             for z in range(FLOORS)]
        history.append((remaining, floor_distribution))
        
        if not agents:
            break
    
    return history

# ==================== 可视化模块 ====================
def visualize_simulation(environment, simulation_data):
    # 创建可视化画布
    fig = plt.figure(figsize=(20, 12))
    grid_spec = fig.add_gridspec(2, 6)
    
    # 3D视图设置
    ax3d = fig.add_subplot(grid_spec[:, 0], projection='3d')
    ax3d.set_xlim(0, WIDTH)
    ax3d.set_ylim(0, HEIGHT)
    ax3d.set_zlim(0, FLOORS)
    ax3d.set_title("三维楼梯动态")
    
    # 绘制楼梯结构
    stair_voxels = (environment == 2)
    ax3d.voxels(stair_voxels, facecolors='sienna', edgecolor='k', alpha=0.7)
    
    # 初始化人员位置
    scatter_3d = ax3d.scatter([], [], [], s=30, c='dodgerblue')
    info_text = ax3d.text2D(0.05, 0.95, "", transform=ax3d.transAxes)
    
    # 各楼层2D视图
    floor_axes = []
    for z in range(FLOORS):
        row = 0 if z < 5 else 1
        col = (z % 5) + 1
        ax = fig.add_subplot(grid_spec[row, col])
        ax.imshow(environment[z], cmap=ListedColormap(['gray', 'white', 'sienna']), 
                origin='lower')
        ax.set_title(f"楼层 {z}")
        ax.axis('off')
        floor_axes.append(ax)
    
    # 动画更新函数
    def update(frame):
        # 更新3D视图
        positions = [p for _, p in simulation_data[frame]]
        if positions:
            xs, ys, zs = zip(*[(x, y, z) for (z, y, x) in positions])
            scatter_3d._offsets3d = (xs, ys, zs)
        else:
            scatter_3d._offsets3d = ([], [], [])
        
        # 更新统计信息
        total, per_floor = simulation_data[frame]
        info_text.set_text(f"剩余人数: {total}\n" + 
                          "\n".join(f"楼层{z}: {per_floor[z]}" for z in range(FLOORS)))
        
        # 更新各楼层视图
        for z in range(FLOORS):
            floor_pos = [(y, x) for (fz, y, x) in positions if fz == z]
            if floor_pos:
                ys, xs = zip(*floor_pos)
                floor_axes[z].collections[1].set_offsets(np.column_stack([xs, ys]))
            else:
                floor_axes[z].collections[1].set_offsets(np.empty((0, 2)))
        
        return [scatter_3d, info_text] + [ax.collections[1] for ax in floor_axes]
    
    # 创建动画
    anim = animation.FuncAnimation(
        fig, update, 
        frames=len(simulation_data), 
        interval=300, 
        blit=False
    )
    
    plt.tight_layout()
    plt.show()

# ==================== 主程序 ====================
if __name__ == "__main__":
    # 生成建筑环境
    building_env = build_environment()
    
    # 计算静态势场
    print("正在计算导航场...")
    navigation_field = compute_static_field(building_env)
    
    # 运行疏散仿真
    print("开始疏散仿真...")
    simulation_history = run_simulation(building_env, navigation_field)
    
    # 可视化结果
    print("生成可视化...")
    visualize_simulation(building_env, simulation_history)